# coding=utf-8
# 学习单位  : 郑州大学
# @Author  : 铭同学
# @Time    : 2021/10/13 21:56
# @Software: PyCharm

# 卷积神经网络CNN(需要弄清楚每个步骤的维度传递)--直接使用cpu，算力直接拉满
# 多分类 Softmax
import torch
from torchvision import transforms       # 图像处理
from torchvision import datasets         # 数据集
from torch.utils.data import DataLoader  # 数据载入
import torch.nn.functional as F          # 函数包--激活函数使用（relu()）
import torch.optim as optim              # 优化器

# 1. Prepare Dataset
batch_size = 64
transform = transforms.Compose([
    transforms.ToTensor(),  # 将PIL Image转换为Tensor（这时通道在前，通道(c)*宽度(w)*高度(h)）
    transforms.Normalize((0.1307, ), (0.3081, ))    # 对数据标准化转换为0-1分布
])

# 准备数据集
train_dataset = datasets.MNIST(root='../dataset/mnist/' #数据路径
                               ,train=True  # 作为训练集
                               ,download=True   # 如果路径无数据进行下载（第一次直接下载）
                               ,transform=transform
                               )
# 批量载入数据集
train_loader = DataLoader(train_dataset #数据集
                          ,shuffle=True # 打乱数据
                          ,batch_size=batch_size    #批量大小
                          )

# 准备数据集
test_dataset = datasets.MNIST(root='../dataset/mnist/' #数据路径
                              ,train=False
                              ,download=True   # 如果路径无数据进行下载
                              ,transform=transform
                              )
# 批量载入数据集
test_loader = DataLoader(test_dataset #数据集
                         ,shuffle=False # 不打乱数据，有利于结果观察
                         ,batch_size=batch_size    #批量大小
                         )
# 2. Design Model
# 模型类--卷积神经网络CNN(需要弄清楚每个步骤的维度)
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, # 输入通道数
                                     10,# 输出通道数
                                     kernel_size=5  # 卷积核5*5矩阵--宽高直接减4(-(n-1))
                                     )  # 卷积层
        # 上一层的输出=下一层的输入
        self.conv2 = torch.nn.Conv2d(10, 20, kernel_size=5) # 卷积层
        self.pooling = torch.nn.MaxPool2d(2)    #池化层--2就是宽高直接砍半（/2）
        self.fc = torch.nn.Linear(320, 10)
    def forward(self, x):
        # Flatten data from (n, 1, 28, 28) to (n, 784)
        batch_size = x.size(0)
        # 激活函数
        x = F.relu(self.pooling(self.conv1(x)))
        x = F.relu(self.pooling(self.conv2(x)))
        # 转换数据维度
        x = x.view(batch_size, -1) # flatten
        # 最后输出线性结果
        x = self.fc(x)
        return x

# 实例化模型
model = Net()

# 迁移到GPU "cuda:0"代表第一块显卡
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# model.to(device)

# 3. Construct Loss and Optimizer
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)    # 加一个动（冲）量（其作用就是摆脱局部最优，奔向全局最优）

# 4. Train and Test
def train(epoch):
    running_loss = 0.0
    # TypeError: __array__() takes 1 positional argument but 2 were given
    # 解决办法，降低pillow版本到8.2即可，pip install pillow==8.2
    for batch_idx, data in enumerate(train_loader, 0): # 这一行读数据失败--已经解决
        inputs, target = data
        # 迁移到GPU
        # inputs, target = inputs.to(device), target.to(device)
        optimizer.zero_grad()   # 梯度归零，避免累加

        # forward + backward + update
        outputs = model(inputs)
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() # item()用于取出值本身，对每个批次的损失进行求和
        # 每300个批次输出一次
        if batch_idx % 300 == 299:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 300))
            running_loss = 0.0 # 损失清零

def test():
    correct = 0
    total = 0
    # 测试时不计算梯度（求梯度是很贵的事情，非必要即可不求）
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            # 迁移到GPU
            # images, labels = inputs.to(device), target.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, dim=1)   # 按照每一行求最大值（就是最大概率），输出两个结果（数值本身和对应下标）我们只要下标predicted(也就是对应的分类)
            total += labels.size(0) # 计算总数
            # 计算预测正确的个数
            correct += (predicted == labels).sum().item()

    print('Accuracy on test set: %.2f%% [%d / %d]' % ((100 * correct / total), correct, total))


if __name__ == '__main__':
    pass
    for epoch in range(10):
        train(epoch)
        test()


''' 训练结果
[1,   300] loss: 0.629
[1,   600] loss: 0.189
[1,   900] loss: 0.142
Accuracy on test set: 97 % [9711 / 10000]
[2,   300] loss: 0.111
[2,   600] loss: 0.100
[2,   900] loss: 0.088
Accuracy on test set: 97 % [9750 / 10000]
[3,   300] loss: 0.077
[3,   600] loss: 0.077
[3,   900] loss: 0.071
Accuracy on test set: 98 % [9834 / 10000]
[4,   300] loss: 0.066
[4,   600] loss: 0.063
[4,   900] loss: 0.060
Accuracy on test set: 98 % [9822 / 10000]
[5,   300] loss: 0.059
[5,   600] loss: 0.051
[5,   900] loss: 0.058
Accuracy on test set: 98 % [9836 / 10000]
[6,   300] loss: 0.051
[6,   600] loss: 0.052
[6,   900] loss: 0.050
Accuracy on test set: 98 % [9866 / 10000]
[7,   300] loss: 0.047
[7,   600] loss: 0.047
[7,   900] loss: 0.044
Accuracy on test set: 98 % [9858 / 10000]
[8,   300] loss: 0.041
[8,   600] loss: 0.043
[8,   900] loss: 0.046
Accuracy on test set: 98 % [9866 / 10000]
[9,   300] loss: 0.040
[9,   600] loss: 0.037
[9,   900] loss: 0.042
Accuracy on test set: 98 % [9872 / 10000]
[10,   300] loss: 0.034
[10,   600] loss: 0.035
[10,   900] loss: 0.043
Accuracy on test set: 98 % [9869 / 10000]
'''